import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime, timedelta
from collections import defaultdict
from abs.DataServiceInterface import IDataService
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import plotly.express as px
class AdvancedBackTrader:

    def __init__(self, df: pd.DataFrame, cash: float = 100000.0,
                 fee_base: float = 1.0,
                 fee_rate: float = 0.0001):
        # 账户初始设置
        self.initial_cash = cash
        self.cash = cash
        self.df = df  # 所有标的的每日收盘价
        self.fee_base = fee_base
        self.fee_rate = fee_rate
        self.transfer_fee_min = 2  # 过户费，只有卖出时收取
        self.transfer_fee_rate = 0.00002
        self.positions = defaultdict(float)  # 持仓数量
        self.trades = defaultdict(list)      # 交易记录
        # 风险指标
        self.max_drawdown = 0

    def buy(self, symbol: str, time: datetime, price: float,
            cash_pct: float = None, amount: int = None):
        """增强买入函数，支持按现金比例或固定数量买入"""
                # 参数类型验证
        if not isinstance(price, (int, float)) or price <= 0:
            raise ValueError(f"无效价格：{price}，必须是正数")
        
        if (cash_pct is None and amount is None) or (cash_pct is not None and amount is not None):
            raise ValueError("必须指定cash_pct或amount中的一个参数")
        
        # 计算买入数量
        if cash_pct is not None:
            if not 0 < cash_pct <= 1:
                raise ValueError("现金比例必须在(0,1]区间")
            
            available_cash = self.cash * cash_pct
            # 计算最大可买手数（1手=100股）
            cost_per_lot = 100 * price
            fee_per_lot = max(self.fee_base, cost_per_lot * self.fee_rate)
            max_lots = int(available_cash // (cost_per_lot + fee_per_lot))
            
            if max_lots < 1:
                raise ValueError(f"可用资金不足买入1手（100股），当前可用：{available_cash:.2f}")
            
            amount = max_lots * 100
        else:
            if amount % 100 != 0:
                raise ValueError("买入数量必须是100的整数倍")
            if amount <= 0:
                raise ValueError("买入数量必须大于0")

        # 计算实际交易金额和费用
        trade_value = amount * price
        fee = max(self.fee_base, trade_value * self.fee_rate)
        
        # 资金检查
        if self.cash < trade_value + fee:
            print(f"资金不足，需要{trade_value+fee:.2f}，可用资金{self.cash:.2f}")
            return 
        
        # 更新账户
        self.cash -= (trade_value + fee)
        self.positions[symbol] += amount
        
        # 记录交易
        self.trades[time].append([symbol, 'buy', price, amount, fee])

    def sell(self, symbol: str, time: datetime, price: float,
             position_pct: float = None, amount: int = None):
        """增强卖出函数，支持按持仓比例或固定数量卖出"""
        # 参数类型验证
        if not isinstance(price, (int, float)) or price <= 0:
            raise ValueError(f"无效价格：{price}，必须是正数")
        
        if (position_pct is None and amount is None) or (position_pct is not None and amount is not None):
            raise ValueError("必须指定position_pct或amount中的一个参数")
        
        current_position = self.positions.get(symbol, 0)
        if current_position < 100:
            raise ValueError(f"持仓不足100股，无法卖出：{symbol}")
        
        # 计算卖出数量
        if position_pct is not None:
            if not 0 < position_pct <= 1:
                raise ValueError("持仓比例必须在(0,1]区间")
            
            # 计算理论卖出数量并向下取整到最近的手数
            desired_shares = current_position * position_pct
            sell_lots = int(desired_shares // 100)
            
            if sell_lots < 1:
                # 至少卖出1手
                sell_lots = 1
            
            amount = sell_lots * 100
            if amount > current_position:
                # 调整到最大可卖数量
                amount = (current_position // 100) * 100
        else:
            if amount % 100 != 0:
                raise ValueError("卖出数量必须是100的整数倍")
            if amount <= 0 or amount > current_position:
                raise ValueError(f"卖出数量无效，当前持仓：{current_position}")

        # 计算费用
        trade_value = amount * price
        fee = self.fee_base
        transfer_fee = max(self.transfer_fee_min, trade_value * self.transfer_fee_rate)
        total_fee = fee + transfer_fee
        
        # 更新账户
        self.cash += (trade_value - total_fee)
        self.positions[symbol] -= amount
        
        # 清理空持仓
        if self.positions[symbol] == 0:
            del self.positions[symbol]
        
        # 记录交易
        self.trades[time].append([symbol, 'sell', price, amount, total_fee])

    def analyze_performance(self, mgr: IDataService, save_to_db=False, backtest_name=None):
        """绩效分析 - 改进版"""
        # 初始化数据结构
        dates = self.df.index.sort_values()
        if len(dates) == 0:
            return {"error": "No data available"}
        
        current_cash = self.initial_cash
        current_positions = defaultdict(float)
        position_cost = defaultdict(list)  # 存储持仓成本 [ (price, amount) ]
        daily_data = []
        total_fees = 0
        trade_dates = set()
        trade_details = []
        
        # 预处理交易记录
        sorted_times = sorted(self.trades.keys())
        all_trades = []
        for time in sorted_times:
            all_trades.extend([(time, trade) for trade in self.trades[time]])
        
        trade_idx = 0
        # 遍历每个交易日
        for date in dates:
            # 处理当日交易
            while trade_idx < len(all_trades):
                trade_time, trade = all_trades[trade_idx]
                if trade_time > date:
                    break
                    
                symbol, action, price, amount, fee = trade
                total_fees += fee
                trade_record = {
                    'date': trade_time,
                    'symbol': symbol,
                    'action': action,
                    'price': price,
                    'amount': amount,
                    'fee': fee
                }
                
                if action == 'buy':
                    current_cash -= (price * amount + fee)
                    current_positions[symbol] += amount
                    # 记录买入成本
                    position_cost[symbol].append((price, amount))
                    trade_record['value'] = -(price * amount + fee)
                    
                elif action == 'sell':
                    current_cash += (price * amount - fee)
                    current_positions[symbol] -= amount
                    trade_record['value'] = price * amount - fee
                    
                    # FIFO计算卖出盈亏
                    remaining = amount
                    total_cost = 0
                    while remaining > 0 and position_cost[symbol]:
                        cost_price, cost_amount = position_cost[symbol][0]
                        if cost_amount <= remaining:
                            total_cost += cost_price * cost_amount
                            remaining -= cost_amount
                            position_cost[symbol].pop(0)
                        else:
                            total_cost += cost_price * remaining
                            position_cost[symbol][0] = (cost_price, cost_amount - remaining)
                            remaining = 0
                    
                    # 计算盈亏
                    profit = (price * amount - fee) - total_cost
                    trade_record['profit'] = profit
                    trade_record['return_pct'] = (profit / total_cost) * 100 if total_cost > 0 else 0
                
                trade_details.append(trade_record)
                trade_dates.add(trade_time)
                trade_idx += 1
            
            # 清理空持仓
            for symbol in list(current_positions.keys()):
                if abs(current_positions[symbol]) < 1e-6:  # 考虑浮点误差
                    del current_positions[symbol]
                    if symbol in position_cost:
                        del position_cost[symbol]
            
            # 计算当日持仓市值
            position_value = 0.0
            for symbol, amount in current_positions.items():
                if symbol in self.df.columns:
                    price = self.df.loc[date, symbol]
                    position_value += amount * price
            
            total_asset = current_cash + position_value
            daily_data.append({
                'date': date,
                'cash': current_cash,
                'position': position_value,
                'total': total_asset
            })
        
        # 创建DataFrame时确保索引是datetime类型
        daily_df = pd.DataFrame(daily_data)
        daily_df['date'] = pd.to_datetime(daily_df['date'])  # 确保转换为datetime
        daily_df = daily_df.set_index('date')  # 设置为索引
        
        # 计算每日收益率
        daily_df['daily_return'] = daily_df['total'].pct_change().fillna(0)
        
        # 计算绩效指标
        final_value = daily_df['total'].iloc[-1]
        total_return = (final_value / self.initial_cash - 1) * 100
        
        # 年化收益率
        days = len(daily_df)
        annual_return = 0
        if days > 0:
            annualized = (final_value / self.initial_cash) ** (252 / days) - 1
            annual_return = annualized * 100
        
        # 最大回撤
        max_peak = daily_df['total'].cummax()
        drawdown = (max_peak - daily_df['total']) / max_peak
        max_drawdown = drawdown.max() * 100
        max_dd_end = drawdown.idxmax()
        max_peak_before = daily_df['total'][:max_dd_end].idxmax()
        max_dd_duration = (max_dd_end - max_peak_before).days
        
        # 计算波动率
        volatility = daily_df['daily_return'].std() * np.sqrt(252) * 100
        
        # 夏普比率
        risk_free_rate = 0.03 / 252
        excess_returns = daily_df['daily_return'] - risk_free_rate
        sharpe_ratio = excess_returns.mean() / excess_returns.std() * np.sqrt(252)
        
        # 索提诺比率
        downside_returns = daily_df['daily_return'].copy()
        downside_returns[downside_returns > 0] = 0
        downside_volatility = downside_returns.std() * np.sqrt(252)
        sortino_ratio = excess_returns.mean() / downside_volatility * np.sqrt(252) if downside_volatility > 0 else 0
        
        # 交易统计
        trade_days = len(trade_dates)
        total_trades = len(all_trades)
        
        # 交易胜率统计
        trade_details_df = pd.DataFrame(trade_details)
        win_trades = loss_trades = 0
        avg_win = avg_loss = 0
        
        if 'profit' in trade_details_df.columns:
            profitable = trade_details_df[trade_details_df['action'] == 'sell']
            win_trades = len(profitable[profitable['profit'] > 0])
            loss_trades = len(profitable[profitable['profit'] < 0])
            
            if win_trades > 0:
                avg_win = profitable[profitable['profit'] > 0]['profit'].mean()
            if loss_trades > 0:
                avg_loss = profitable[profitable['profit'] < 0]['profit'].mean()
        
        win_rate = win_trades / total_trades * 100 if total_trades > 0 else 0
        profit_factor = abs(avg_win * win_trades / (avg_loss * loss_trades)) if loss_trades > 0 else float('inf')
        
        # 使用英文键名生成报告
        report = {
            'name': backtest_name or datetime.now().strftime("%Y%m%d_%H%M%S"),
            'start_time': self.df.index.min(),
            'end_time': self.df.index.max(),
            'initial_cash': self.initial_cash,
            'final_value': final_value,
            'total_return': total_return,  # 百分比值
            'annual_return': annual_return,
            'max_drawdown': max_drawdown,
            'total_fees': total_fees,
            'trade_days': trade_days,
            'total_trades': total_trades,
            'sharpe_ratio': sharpe_ratio,
            'sortino_ratio': sortino_ratio,
            'win_trades': win_trades,
            'loss_trades': loss_trades,
            'win_rate': win_rate,
            'avg_win': avg_win,
            'avg_loss': avg_loss,
            'profit_factor': profit_factor,
            'max_drawdown_start': max_peak_before,
            'max_drawdown_end': max_dd_end,
            'max_drawdown_duration': max_dd_duration,
            'volatility': volatility,
            'end_position_count': len(current_positions),
            'end_position_ratio': (position_value / total_asset) * 100 if total_asset > 0 else 0
        }
        
        # 保存到数据库
        if save_to_db and mgr is not None:
            mgr.save_backtest_result(report, daily_df[['cash', 'position', 'total','daily_return']],trade_details_df)
        
        # # 生成图表
        # self._plot_performance(daily_df, trade_details_df)
        
        return report

    def ___plot_performance(self, daily_df: pd.DataFrame, trade_details_df: pd.DataFrame = None):
        """
        已经弃用，使用的plt，功能没有问题，但是与前端不兼容
        生成可视化图表 - 增强版"""
        fig = plt.figure(figsize=(16, 14))
        
        # 资金曲线
        plt.subplot(4, 2, 1)
        plt.plot(daily_df['total'], label='总资产')
        plt.plot(daily_df['cash'], label='现金')
        plt.plot(daily_df['position'], label='持仓市值')
        plt.title('资金曲线')
        plt.grid(True)
        plt.legend()
        
        # 回撤曲线
        plt.subplot(4, 2, 2)
        max_peak = daily_df['total'].cummax()
        drawdown = (max_peak - daily_df['total']) / max_peak
        plt.fill_between(daily_df.index, drawdown * 100, 0, color='red', alpha=0.3)
        plt.title('回撤曲线')
        plt.ylabel('回撤 (%)')
        plt.grid(True)
        
        # 交易频次（如果有交易细节）
        if trade_details_df is not None and not trade_details_df.empty:
            plt.subplot(4, 2, 3)
            trade_details_df['date'] = pd.to_datetime(trade_details_df['date'])
            trade_counts = trade_details_df.groupby(trade_details_df['date'].dt.date).size()
            trade_counts.plot(kind='bar', width=0.8)
            plt.title('交易频次')
            plt.xlabel('日期')
            plt.ylabel('交易次数')
        else:
            # 如果没有交易细节，留空或者显示提示
            plt.subplot(4, 2, 3)
            plt.text(0.5, 0.5, '无交易记录', ha='center', va='center')
            plt.title('交易频次')
        
        # 收益率分布
        plt.subplot(4, 2, 4)
        if 'daily_return' in daily_df.columns:
            daily_df['daily_return'].hist(bins=50, alpha=0.7)
            plt.title('日收益率分布')
            plt.xlabel('日收益率')
            plt.ylabel('频率')
        
        # 累积收益率
        plt.subplot(4, 2, 5)
        if 'daily_return' in daily_df.columns:
            cumulative_return = (1 + daily_df['daily_return']).cumprod() - 1
            cumulative_return.plot()
            plt.title('累积收益率')
            plt.ylabel('收益率')
            plt.grid(True)
        
        # 持仓比例
        plt.subplot(4, 2, 6)
        position_ratio = daily_df['position'] / daily_df['total']
        position_ratio.plot()
        plt.title('持仓比例')
        plt.ylabel('比例')
        plt.ylim(0, 1)
        plt.grid(True)
        
        # 盈利交易分析（如果有交易细节）
        if trade_details_df is not None and not trade_details_df.empty and 'profit' in trade_details_df.columns:
            plt.subplot(4, 2, 7)
            win_trades = trade_details_df[trade_details_df['profit'] > 0]
            loss_trades = trade_details_df[trade_details_df['profit'] < 0]
            
            plt.hist(win_trades['profit'], bins=30, alpha=0.7, color='green', label='盈利')
            plt.hist(loss_trades['profit'], bins=30, alpha=0.7, color='red', label='亏损')
            plt.title('交易盈亏分布')
            plt.xlabel('盈亏金额')
            plt.ylabel('交易次数')
            plt.legend()
        else:
            plt.subplot(4, 2, 7)
            plt.text(0.5, 0.5, '无交易盈亏数据', ha='center', va='center')
            plt.title('交易盈亏分布')
        
        plt.tight_layout()
        
        return fig  # 返回图表对象而不是直接显示

    def _plot_performance(self, daily_df: pd.DataFrame, trade_details_df: pd.DataFrame = None):
        """使用Plotly生成交互式图表 - 增强版"""

        fig = make_subplots(
            rows=4, cols=2,
            subplot_titles=('资金曲线', '回撤曲线', '交易频次', '日收益率分布', 
                           '累积收益率', '持仓比例', '交易盈亏分布'),
            specs=[
                [{"type": "scatter"}, {"type": "scatter"}],
                [{"type": "bar"}, {"type": "histogram"}],
                [{"type": "scatter"}, {"type": "scatter"}],
                [{"type": "histogram"}, {}]
            ],
            vertical_spacing=0.08,
            horizontal_spacing=0.1
        )
        # 修复时间轴问题 - 确保日期格式正确
        daily_df.index = pd.to_datetime(daily_df.index)
        if not daily_df.empty:
            # 设置图表的时间范围从回测开始前一天到结束后一天
            start_date = daily_df.index.min() - pd.Timedelta(days=1)
            end_date = daily_df.index.max() + pd.Timedelta(days=1)

        # 更新图表的时间范围
        if 'xaxis' in fig.layout:
            fig.update_layout(
                xaxis=dict(range=[start_date, end_date]),
                xaxis2=dict(range=[start_date, end_date]),
                xaxis5=dict(range=[start_date, end_date]),
                xaxis6=dict(range=[start_date, end_date])
            )
        # 资金曲线
        fig.add_trace(go.Scatter(
            x=daily_df.index, 
            y=daily_df['total'],
            name='总资产',
            line=dict(color='royalblue', width=2),
            hovertemplate='日期: %{x|%Y-%m-%d}<br>总资产: %{y:,.2f}元<extra></extra>'
        ), row=1, col=1)
        
        fig.add_trace(go.Scatter(
            x=daily_df.index, 
            y=daily_df['cash'],
            name='现金',
            line=dict(color='green', width=2),
            hovertemplate='日期: %{x|%Y-%m-%d}<br>现金: %{y:,.2f}元<extra></extra>'
        ), row=1, col=1)
        
        fig.add_trace(go.Scatter(
            x=daily_df.index, 
            y=daily_df['position'],
            name='持仓市值',
            line=dict(color='orange', width=2),
            hovertemplate='日期: %{x|%Y-%m-%d}<br>持仓市值: %{y:,.2f}元<extra></extra>'
        ), row=1, col=1)
        
        # 回撤曲线
        max_peak = daily_df['total'].cummax()
        drawdown = (max_peak - daily_df['total']) / max_peak * 100
        
        fig.add_trace(go.Scatter(
            x=daily_df.index,
            y=drawdown,
            fill='tozeroy',
            fillcolor='rgba(255, 0, 0, 0.3)',
            line=dict(color='red', width=1),
            name='回撤',
            hovertemplate='日期: %{x|%Y-%m-%d}<br>回撤: %{y:.2f}%<extra></extra>'
        ), row=1, col=2)
        
        # 交易频次
        if trade_details_df is not None and not trade_details_df.empty:
            trade_details_df['date'] = pd.to_datetime(trade_details_df['date'])
            trade_counts = trade_details_df.groupby(trade_details_df['date'].dt.date).size()
            
            fig.add_trace(go.Bar(
                x=trade_counts.index,
                y=trade_counts.values,
                name='交易次数',
                marker_color='cornflowerblue',
                hovertemplate='日期: %{x}<br>交易次数: %{y}次<extra></extra>'
            ), row=2, col=1)
        else:
            fig.add_annotation(
                text="无交易记录",
                xref="x domain", yref="y domain",
                x=0.5, y=0.5, showarrow=False,
                row=2, col=1
            )
        
        # 日收益率分布
        if 'daily_return' in daily_df.columns:
            fig.add_trace(go.Histogram(
                x=daily_df['daily_return'] * 100,
                nbinsx=50,
                name='收益率分布',
                marker_color='lightseagreen',
                hovertemplate='收益率: %{x:.2f}%<br>频率: %{y}天<extra></extra>'
            ), row=2, col=2)
        
        # 累积收益率
        if 'daily_return' in daily_df.columns:
            cumulative_return = (1 + daily_df['daily_return']).cumprod() - 1
            
            fig.add_trace(go.Scatter(
                x=daily_df.index,
                y=cumulative_return * 100,
                name='累积收益',
                line=dict(color='purple', width=2),
                hovertemplate='日期: %{x|%Y-%m-%d}<br>收益率: %{y:.2f}%<extra></extra>'
            ), row=3, col=1)
        
        # 持仓比例
        position_ratio = daily_df['position'] / daily_df['total']
        
        fig.add_trace(go.Scatter(
            x=daily_df.index,
            y=position_ratio * 100,
            name='持仓比例',
            line=dict(color='teal', width=2),
            hovertemplate='日期: %{x|%Y-%m-%d}<br>持仓比例: %{y:.2f}%<extra></extra>'
        ), row=3, col=2)
        
        # 交易盈亏分布
        if trade_details_df is not None and not trade_details_df.empty and 'profit' in trade_details_df.columns:
            win_trades = trade_details_df[trade_details_df['profit'] > 0]
            loss_trades = trade_details_df[trade_details_df['profit'] < 0]
            
            fig.add_trace(go.Histogram(
                x=win_trades['profit'],
                name='盈利交易',
                marker_color='forestgreen',
                opacity=0.7,
                hovertemplate='盈利: %{x:.2f}元<br>次数: %{y}次<extra></extra>'
            ), row=4, col=1)
            
            fig.add_trace(go.Histogram(
                x=loss_trades['profit'],
                name='亏损交易',
                marker_color='firebrick',
                opacity=0.7,
                hovertemplate='亏损: %{x:.2f}元<br>次数: %{y}次<extra></extra>'
            ), row=4, col=1)
        else:
            fig.add_annotation(
                text="无交易盈亏数据",
                xref="x domain", yref="y domain",
                x=0.5, y=0.5, showarrow=False,
                row=4, col=1
            )
        
        # 更新图表布局
        fig.update_layout(
            title='回测绩效分析',
            height=1200,
            showlegend=True,
            hovermode='x unified',
            template='plotly_white',
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            )
        )
        
        # 更新子图标题和坐标轴标签
        fig.update_yaxes(title_text="金额(元)", row=1, col=1)
        fig.update_yaxes(title_text="回撤(%)", row=1, col=2)
        fig.update_yaxes(title_text="交易次数", row=2, col=1)
        fig.update_yaxes(title_text="频率", row=2, col=2)
        fig.update_yaxes(title_text="收益率(%)", row=3, col=1)
        fig.update_yaxes(title_text="持仓比例(%)", row=3, col=2)
        fig.update_yaxes(title_text="交易次数", row=4, col=1)
        
        fig.update_xaxes(title_text="日期", row=1, col=1)
        fig.update_xaxes(title_text="日期", row=1, col=2)
        fig.update_xaxes(title_text="日期", row=2, col=1)
        fig.update_xaxes(title_text="收益率(%)", row=2, col=2)
        fig.update_xaxes(title_text="日期", row=3, col=1)
        fig.update_xaxes(title_text="日期", row=3, col=2)
        fig.update_xaxes(title_text="盈亏金额(元)", row=4, col=1)
        
        return fig

    @classmethod
    def load_and_show_backtest(cls, backtest_name, mgr:IDataService):
        """从数据库加载回测结果并展示（支持交易细节）"""
        # 从数据库获取回测结果（现在返回三个值）
        result = mgr.get_backtest_detail(backtest_name)
        if result is None or result[0] is None:
            print(f"回测 '{backtest_name}' 不存在")
            return
        
        summary, daily_df, trade_details_df = result
        
        # 创建回测实例
        bt = cls(pd.DataFrame(), cash=summary['initial_cash'])
        
        # # 绘制结果
        # bt._plot_performance(daily_df, trade_details_df)
        
        return summary, daily_df, trade_details_df

    @classmethod
    def list_backtests(cls, mgr:IDataService):
        """列出所有回测结果"""
        backtests = mgr.get_backtest_list()
        if not backtests:
            print("没有找到任何回测结果")
            return
        
        print("\n" + "="*50)
        print("回测结果列表:")
        for i, bt in enumerate(backtests, 1):
            print(f"{i}. {bt['name']} ({bt['start_time'].date()} 至 {bt['end_time'].date()})")
        print("="*50 + "\n")
        
        return backtests